Goto

Collaborating Authors

 nuisance function


Learning heterogeneous treatment effects under principal stratification

arXiv.org Machine Learning

Principal stratification provides a foundational framework for causal inference with intermediate outcomes by defining causal effects within subpopulations, yet existing work has largely focused on average effects across strata rather than treatment effect heterogeneity within strata. Such within-stratum heterogeneity informs individualized treatment decisions but the associated methods are sparse. We address this gap by studying the identification and estimation of the conditional principal causal effects under principal ignorability combined with an odds ratio sensitivity parameterization, which relaxes the monotonicity assumption. To efficiently learn these estimands, we propose a novel doubly cross-fit doubly robust machine learner that resolves the nested nuisance structure inherent to principal stratification. Leveraging sequential orthogonal learning with regularized least-squares sieves, we derive $\mathcal{L}^2$ and uniform limit theory, establish oracle efficiency, and construct uniform confidence bands for the proposed estimator. We use simulations to demonstrate the finite-sample performance of our estimator, and provide an empirical analysis of a randomized trial in acute lung injury, revealing informative patterns of treatment effect heterogeneity within the always-survivor subpopulation.


f0156a82b6af6a4e838923ce9c124424-Paper-Conference.pdf

Neural Information Processing Systems

Structure-agnostic causal inference studies how well one can estimate a treatment effect given black-box machine learning estimates of nuisance functions (like the impact of confounders on treatment and outcomes). Here, we find that the answer depends in a surprising way on the distribution of the treatment noise. Focusing on the partially linear model of Robinson [1988], we first show that the widely adopted double machine learning (DML) estimator is minimax rate-optimal for Gaussian treatment noise, resolving an open problem of Mackey et al. [2018]. Meanwhile, for independent non-Gaussian treatment noise, we show that DML is always suboptimal by constructing new practical procedures with higher-order robustness to nuisance errors. These ACE procedures use structure-agnostic cumulant estimators to achieve r-th order insensitivity to nuisance errors whenever the (r + 1)-st treatment cumulant is non-zero. We complement these core results with novel minimax guarantees for binary treatments in the partially linear model. Finally, using synthetic demand estimation experiments, we demonstrate the practical benefits of our higher-order robust estimators.


Treatment Effect Estimation for Optimal Decision-Making

Neural Information Processing Systems

Decision-making in various fields, such as medicine, is heavily based on conditional average treatment effects (CATEs). Practitioners commonly make decisions by checking whether the estimated CATE is positive, even though the decision-making performance of modern CATE estimators (meta-learners) is poorly understood. In this paper, we study optimal decision-making based on two-stage meta-learners (e.g., DR-learner), which estimate CATE via a second-stage regression. We show that these meta-learners can be suboptimal when used for decision-making in common settings where the second-stage regression is over a restricted function class (e.g., when using regularization or employing fairness/interpretability constraints). Intuitively, this occurs because such estimators prioritize CATE accuracy in regions far away from the decision boundary, which is ultimately irrelevant to decision-making. As a remedy, we propose a novel two-stage learning objective that re-targets the CATE to balance CATE estimation error and decision performance. We then propose a neural method that optimizes an adaptively-smoothed approximation of our learning objective. Finally, we confirm the effectiveness of our method both empirically and theoretically.


Optimal Nuisance Function Tuning for Estimating a Doubly Robust Functional under Proportional Asymptotics

Neural Information Processing Systems

In this paper, we explore the asymptotically optimal tuning parameter choice in ridge regression for estimating nuisance functions of a statistical functional that has recently gained prominence in conditional independence testing and causal inference. Given a sample of size n, we study estimators of the Expected Conditional Covariance (ECC) between variables Y and Agiven a high-dimensional covariate X Rp. Under linear regression models for Y and A on X and the proportional asymptotic regime p/n c (0,), we evaluate three existing ECC estimators and two sample splitting strategies for estimating the required nuisance functions. Since no consistent estimator of the nuisance functions exists in the proportional asymptotic regime without imposing further structure on the problem, we first derive debiased versions of the ECC estimators that utilize the ridge regression nuisance function estimators. We show that our bias correction strategy yields n-consistent estimators of the ECC across different sample splitting strategies and estimator choices. We then derive the asymptotic variances of these debiased estimators to illustrate the nuanced interplay between the sample splitting strategy, estimator choice, and tuning parameters of the nuisance function estimators for optimally estimating the ECC. Our analysis reveals that prediction-optimal tuning parameters (i.e., those that optimally estimate the nuisance functions) may not lead to the lowest asymptotic variance of the ECC estimator - thereby demonstrating the need to be careful in selecting tuning parameters based on the final goal of inference. Finally, we verify our theoretical results through extensive numerical experiments.


Orthogonal Survival Learners for Estimating Heterogeneous Treatment Effects from Time-to-Event Data

Neural Information Processing Systems

Estimating heterogeneous treatment effects (HTEs) is crucial for personalized decision-making. However, this task is challenging in survival analysis, which includes time-to-event data with censored outcomes (e.g., due to study dropout). In this paper, we propose a toolbox of orthogonal survival learners to estimate HTEs from time-to-event data under censoring. Our learners have three main advantages: (i) we show that learners from our toolbox are guaranteed to be orthogonal and thus robust with respect to nuisance estimation errors; (ii) our toolbox allows for incorporating a custom weighting function, which can lead to robustness against different types of low overlap, and (iii) our learners are modelagnostic (i.e., they can be combined with arbitrary machine learning models). We instantiate the learners from our toolbox using several weighting functions and, as a result, propose various neural orthogonal survival learners. Some of these coincide with existing survival learners (including survival versions of the DRand R-learner), while others are novel and further robust w.r.t.


Optimally taming biases in black-box models for efficient semiparametric estimation

arXiv.org Machine Learning

Modern semiparametric estimation often relies on flexible black-box machine learning methods to estimate nuisance functions, raising a fundamental question: how do nuisance estimation errors propagate into inference for low-dimensional target parameters? The dominant paradigm, exemplified by double machine learning (DML), yields error bounds in which nuisance estimation errors enter multiplicatively. While widely adopted, it remains unclear whether this multiplicative-rate dependence is optimal for black-box models. In this paper, we start by revisiting the partial linear model $Y = μ_0(X)+T\cdotβ_0+\varepsilon$ under a structure-agnostic setting, where the nuisance function $μ_0$ is estimated using a generic machine learning model, with approximation error $δ^a_μ$ and stochastic error $δ_μ^s$. We show that the standard DML rate is not optimal in the regime where the auxiliary function $\mathbb{E}[T|X=x]$ cannot be consistently estimated. We propose a new estimator for $β_0$ that achieves a sharper rate of $n^{-1/2}+δ^a_μ+(δ_μ^s)^2$ and establish a matching lower bound demonstrating its optimality. Our results reveal a new principle: the first-order stochastic error of nuisance estimation can be eliminated without imposing any additional assumptions. This also leads to a revised tuning strategy favoring under-smoothing, where $δ^a_μ\asymp(δ_μ^s)^2$, rather than the classical bias-variance trade-off $δ^a_μ\asymp δ_μ^s$. Under mild additional conditions, the estimator is asymptotically normal with minimal asymptotic variance. The proposed method extends to a broad class of semi-parametric linear functional estimation problems, including average treatment effect estimation. Our results imply that popular orthogonal score methods in semiparametric estimation with black-box nuisance learners can be substantially improved.


Causal Label Recovery in Payment Networks

arXiv.org Machine Learning

Fraud detection models in payment networks train on chargeback labels that are systematically biased. Every label must survive three sequential gates: authorization (declined transactions generate no labels), issuer reporting (unreported fraud is invisible), and delay (pending chargebacks are missing at training time). Labels that do arrive may be corrupted by first-party misuse or issuer misclassification. A companion paper [arXiv:2605.27557] proved that these four impairments impose a minimax lower bound on detection performance. This paper asks: can that bound be achieved? We formalize the observation pipeline as a sequential missing-data problem with three propensity stages and a corruption layer, and construct the Sequential Triply Robust (STR) estimator. The STR corrects for all four impairments simultaneously and achieves the semiparametric efficiency bound -- no estimator can have lower asymptotic variance. It is sequentially triply robust: at each gate, consistency requires only that either the propensity model or the outcome regression is correctly specified, not both. We provide corruption correction via noise-rate-adjusted pseudo-labels, empirical Bayes shrinkage to stabilize inverse-propensity weights for small issuers, a plug-in variance estimator yielding valid confidence intervals, and a Bernstein concentration inequality for finite-sample guarantees. On the operational side, we derive the optimal training delay -- the maturity window that minimizes the sum of label-quality loss and model staleness -- and prove that the STR permits training on data that is days old rather than months old, decoupling model freshness from the chargeback maturity cycle. The STR provably dominates naive chargeback-based training in mean squared error for any sample size.


Proximal Path-Specific Inference

arXiv.org Machine Learning

Mediation analysis (Robins & Greenland 1992, Pearl 2001, Imai, Keele & Tingley 2010, Tchetgen Tchetgen & Shpitser 2012) provides a principled framework for investigating causal mechanisms by decomposing the effect of a treatment A on an outcome Y into pathways operating through a mediator of interest M. Classical mediation analysis focuses on the natural indirect effect, corresponding to the pathway from Ato Y through M, and the natural direct effect, corresponding to pathways not through M. These estimands are well understood when a single mediator is present and strong identification assumptions hold. However, in many applications, there exist multiple intermediate variables between treatment and outcome. In such settings, conventional mediation analysis typically requires the absence of treatment-induced mediator-outcome confounders--often referred to as recanting witnesses--as well as the absence of unmeasured confounding. Under these circumstances, commonly used identification assumptions such as sequential ignorability (Imai, Keele & Yamamoto 2010) or nonparametric structural equation models with independent errors (NPSEM-IE) (Pearl 2009) no longer suffice to identify natural indirect effects (Avin et al. 2005, Tchetgen Tchetgen & VanderWeele 2014). Figure 1 illustrates this issue: the recanting witness D is directly affected by A and simultaneously confounds the relationship between M and Y. Such treatment-induced confounding is common in epidemiologic studies, particularly when the mediator of interest occurs long after the treatment initiation (Robins 1999). A motivating example arises in studies of preterm birth. Mediation analysis has been widely used to explore whether adequate prenatal care (A) reduces the risk of preterm birth (Y) through preeclampsia (M) (Vansteelandt & VanderWeele 2012, VanderWeele et al. 2014, Xia & Chan 2023).


Fast Uncertainty Quantification for Kernel-Based Estimators in Large-Scale Causal Inference

arXiv.org Machine Learning

Kernel methods are widely used in causal inference for tasks such as treatment effect estimation, policy evaluation, and policy learning. The bootstrap is a standard tool for uncertainty quantification because of its broad applicability. As increasingly large datasets become available, such as the 2023 U.S. Natality data from the National Vital Statistics System (NVSS), which includes 3,596,017 registered births, the computational demands of these methods increase substantially. Kernel methods are known to scale poorly with sample size, and this limitation is further exacerbated by the repeated re-fitting required by the bootstrap. As a result, bootstrap-based inference for kernel-based estimators can become computationally infeasible in large-scale settings. In this paper, we address these challenges by extending the causal Bag of Little Bootstraps (cBLB) algorithm to kernel methods. Our approach achieves computational scalability by combining subsampling and resampling while preserving first-order uncertainty quantification and asymptotically correct coverage. We evaluate the method across three representative implementations: kernelized augmented outcome-weighted learning, kernel-based minimax weighting, and double machine learning with kernel support vector machines. We show in simulations that our method yields confidence intervals with nominal coverage at a fraction of the computational cost. We further demonstrate its utility in a real-world application by estimating the effect of any amount of smoking on birth weight, as well as the optimal treatment regime, using the NVSS dataset, where the standard bootstrap is prohibitively expensive computationally and effectively infeasible at this scale.